package kyber

import (
	"encoding/hex"
	"fmt"
	"test/gincircl/common"

	"github.com/cloudflare/circl/kem"
	"github.com/cloudflare/circl/kem/kyber/kyber1024"
	"github.com/cloudflare/circl/kem/kyber/kyber512"
	"github.com/cloudflare/circl/kem/kyber/kyber768"
)

type KyberWrapper struct {
	scheme kem.Scheme
}

// 创建新的 Kyber 实例
func NewKyber(variant string) (*KyberWrapper, error) {
	switch variant {
	case "512":
		return &KyberWrapper{scheme: kyber512.Scheme()}, nil
	case "768":
		return &KyberWrapper{scheme: kyber768.Scheme()}, nil
	case "1024":
		return &KyberWrapper{scheme: kyber1024.Scheme()}, nil
	default:
		return nil, fmt.Errorf("不支持的 Kyber 变体: %s", variant)
	}
}

// 生成密钥对
func (k *KyberWrapper) GenerateKeyPair() (string, string, error) {
	pk, sk, err := k.scheme.GenerateKeyPair()
	if err != nil {
		common.Logger.Error("生成密钥对失败:" + err.Error())
		return "", "", fmt.Errorf("生成密钥对失败: %v", err)
	}
	pkBytes, err := pk.MarshalBinary()
	if err != nil {
		common.Logger.Error("公钥序列化失败:" + err.Error())
		return "", "", err
	}

	skBytes, err := sk.MarshalBinary()
	if err != nil {
		common.Logger.Error("私钥序列化失败:" + err.Error())
		return "", "", err
	}

	return hex.EncodeToString(pkBytes), hex.EncodeToString(skBytes), nil
}

// 密钥封装
func (k *KyberWrapper) Encapsulate(pkHex string) (string, string, error) {
	pkBytes, err := hex.DecodeString(pkHex)
	if err != nil {
		common.Logger.Error("DecodeString解码公钥失败" + err.Error())
		return "", "", fmt.Errorf("解码公钥失败: %v", err)
	}
	pk, err := k.scheme.UnmarshalBinaryPublicKey(pkBytes)
	if err != nil {
		common.Logger.Error("UnmarshalBinaryPublicKeyg解码公钥失败" + err.Error())
		return "", "", fmt.Errorf("解码公钥失败: %v", err)
	}
	ct, ss, err := k.scheme.Encapsulate(pk)
	if err != nil {
		common.Logger.Error("Encapsulate密钥封装失败" + err.Error())
		return "", "", fmt.Errorf("密钥封装失败: %v", err)
	}
	return hex.EncodeToString(ct), hex.EncodeToString(ss), nil
}

// 密钥解封装
func (k *KyberWrapper) Decapsulate(skHex string, ct string) (string, error) {
	skBytes, err := hex.DecodeString(skHex)
	if err != nil {
		common.Logger.Error("DecodeString解码私钥失败:" + err.Error())
		return "", fmt.Errorf("解码私钥失败: %v", err)
	}
	sk, err := k.scheme.UnmarshalBinaryPrivateKey(skBytes)
	if err != nil {
		common.Logger.Error("UnmarshalBinaryPrivateKey解码私钥失败:" + err.Error())
		return "", fmt.Errorf("密钥解封装失败: %v", err)
	}
	ctBytes, err := hex.DecodeString(ct)
	if err != nil {
		common.Logger.Error("DecodeString解码密文失败:" + err.Error())
		return "", fmt.Errorf("解码密文失败: %v", err)
	}
	ss, err := k.scheme.Decapsulate(sk, ctBytes)
	if err != nil {
		common.Logger.Error("Decapsulate密钥解封装失败:" + err.Error())
		return "", fmt.Errorf("密钥解封装失败: %v", err)
	}
	return hex.EncodeToString(ss), nil
}
